# -*- coding: utf-8 -*-
"""
@author: YuHaiyang

"""
import os
import shutil
from pathlib import Path, PosixPath

import PIL
from PIL.Image import Image
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms


class DataSetLoader:
    def __init__(self, path):
        print("init path:", type(path))
        if isinstance(path, str):
            self.src_path = Path(path)
        elif isinstance(path, PosixPath):
            self.src_path = path

        print("src_path", self.src_path)

    def spilt(self, pattern: str, out_path):
        if isinstance(out_path, str):
            out_path = Path(out_path)
        elif isinstance(out_path, PosixPath):
            out_path = out_path

        if not out_path.exists():
            out_path.mkdir(exist_ok=True)

        for i in self.src_path.glob(pattern):
            shutil.copy2(i, Path(out_path, i.name))

    def gen(self, batch_size: int = 256, num_workers: int = 0) -> DataLoader:
        for ignore in self.src_path.rglob(".DS_Store"):
            print("ignore:", ignore.absolute())
            os.remove(ignore)

        init_trans = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        data_set = ImageFolder(root=str(self.src_path.absolute()),
                               transform=init_trans,
                               is_valid_file=lambda x: not x.endswith('.DS_Store')
                               )

        print("data_set len:", len(data_set))
        print("class_to_idx len:", data_set.class_to_idx)

        data_loader = DataLoader(data_set, batch_size=batch_size, shuffle=True, num_workers=num_workers)

        return data_loader

    def gen_src_imgs(self, out_path: str, count: int, enhancement: bool = True, size: int = 65):
        out_path = Path(out_path)
        if not out_path.exists():
            out_path.mkdir(exist_ok=True)

        for i in self.src_path.iterdir():
            print("i:", i)
            img = PIL.Image.open(i)

            transform = transforms.Compose([
                transforms.ToTensor(),
                transforms.ToPILImage(),
                transforms.Resize(size),
            ])

            img = transform(img)
